# ------------------------------------------------------------------------------
# Get shift coefficients + bins for df
# Author: cshubatt@gmail.com <cshubatt@gmail.com>
# To run: bash 04_get_shift_coefs.sh
# ------------------------------------------------------------------------------

# Libraries --------------------------------------------------------------------
library(here)
library(yaml)
library(data.table)
library(tidyverse)
library(testit) # assert()
library(glue) # glue strings
library(broom) # tidy regression results
library(OneR) # bin

temp <- here::here("code", "05_natural_experiment", "temp")
overnight_lab <- ""
u <- modules::use(here::here("lib", "util.R"))

# Helper Functions -------------------------------------------------------------
get_shift_coefs <- function(model, lab, ...) {
  print(lab)
  shift_var <- str_extract(lab, "shift_\\d\\d")
  df <- copy(shifts)
  coefs <- coef(model)[[shift_var]] %>%
    select(`(Intercept)`) %>%
    rownames_to_column() %>%
    mutate(rowname = factor(as.numeric(rowname))) %>%
    setnames(
      c("rowname", "(Intercept)"),
      c(shift_var, "shift_test_effect")
    )

  df <- df %>%
    u$safe_left_join(coefs) %>%
    filter(!is.na(shift_test_effect)) %>%
    mutate(
      shift_test_bin = bin(
        shift_test_effect,
        nbins = 4, labels = 1:4,
        method = "content"
      )
    )

  df <- df %>%
    select(
      all_of(
        c(
          "ed_enc_id", shift_var, "shift_test_bin", "shift_test_effect"
        )
      )
    ) %>%
    setnames(
      c("shift_test_bin", "shift_test_effect"),
      c(glue("bin_{lab}"), lab)
    )
  assert("df is non-empty", nrow(df) > 0)
  return(df)
}

# Load Data --------------------------------------------------------------------
message("Loading data...")
paths <- read_yaml(here::here("lib", "filepaths.yml"))
models_df <- readRDS(file.path(temp, "re_models.rds"))
# window <- c("12", "08")
window <- c("12")
shift_vars <- glue("shift_{window}")
shifts <- readRDS(file.path(temp, "shift_test_rates.rds"))

# Get Coefficient Bins ---------------------------------------------------------
message("Getting coefficients and bins...")
models_coefs <- models_df %>%
  mutate(coef_bins = pmap(., get_shift_coefs))

# Join Coefs -------------------------------------------------------------------
message("Joining testing coefficient bins...")
coef_bins <- models_coefs$coef_bins
for (i in 1:length(coef_bins)) {
  shifts <- shifts %>%
    u$safe_left_join(coef_bins[[i]])
}

# Save -------------------------------------------------------------------------
save_fp <- file.path(temp, "shift_bins_df.rds")
message("Saving shift RE coefs to ", save_fp, "...")
write_rds(shifts, save_fp)

message("Done.")
